// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


// Adapted from Apache Arrow.

#pragma once

#include <map>
#include <memory>
#include <string>
#include <vector>

#include <pollux/dwio/parquet/writer/arrow/schema.h>

namespace kumo::pollux::parquet::arrow {
    namespace encryption {
        class AesDecryptor;
        class AesEncryptor;
    } // namespace encryption

    class FileDecryptionProperties;

    class PARQUET_EXPORT Decryptor {
    public:
        Decryptor(
            std::shared_ptr<encryption::AesDecryptor> decryptor,
            const std::string &key,
            const std::string &file_aad,
            const std::string &aad,
            ::arrow::MemoryPool *pool);

        const std::string &file_aad() const {
            return file_aad_;
        }

        void UpdateAad(const std::string &aad) {
            aad_ = aad;
        }

        ::arrow::MemoryPool *pool() {
            return pool_;
        }

        int CiphertextSizeDelta();

        int Decrypt(
            const uint8_t *ciphertext,
            int ciphertext_len,
            uint8_t *plaintext);

    private:
        std::shared_ptr<encryption::AesDecryptor> aes_decryptor_;
        std::string key_;
        std::string file_aad_;
        std::string aad_;
        ::arrow::MemoryPool *pool_;
    };

    class InternalFileDecryptor {
    public:
        explicit InternalFileDecryptor(
            FileDecryptionProperties *properties,
            const std::string &file_aad,
            ParquetCipher::type algorithm,
            const std::string &footer_key_metadata,
            ::arrow::MemoryPool *pool);

        std::string &file_aad() {
            return file_aad_;
        }

        std::string GetFooterKey();

        ParquetCipher::type algorithm() {
            return algorithm_;
        }

        std::string &footer_key_metadata() {
            return footer_key_metadata_;
        }

        FileDecryptionProperties *properties() {
            return properties_;
        }

        void WipeOutDecryptionKeys();

        ::arrow::MemoryPool *pool() {
            return pool_;
        }

        std::shared_ptr<Decryptor> GetFooterDecryptor();

        std::shared_ptr<Decryptor> GetFooterDecryptorForColumnMeta(
            const std::string &aad = "");

        std::shared_ptr<Decryptor> GetFooterDecryptorForColumnData(
            const std::string &aad = "");

        std::shared_ptr<Decryptor> GetColumnMetaDecryptor(
            const std::string &column_path,
            const std::string &column_key_metadata,
            const std::string &aad = "");

        std::shared_ptr<Decryptor> GetColumnDataDecryptor(
            const std::string &column_path,
            const std::string &column_key_metadata,
            const std::string &aad = "");

    private:
        FileDecryptionProperties *properties_;
        // Concatenation of aad_prefix (if exists) and aad_file_unique
        std::string file_aad_;
        std::map<std::string, std::shared_ptr<Decryptor> > column_data_map_;
        std::map<std::string, std::shared_ptr<Decryptor> > column_metadata_map_;

        std::shared_ptr<Decryptor> footer_metadata_decryptor_;
        std::shared_ptr<Decryptor> footer_data_decryptor_;
        ParquetCipher::type algorithm_;
        std::string footer_key_metadata_;
        // A weak reference to all decryptors that need to be wiped out when
        // decryption is finished
        std::vector<std::weak_ptr<encryption::AesDecryptor> > all_decryptors_;

        ::arrow::MemoryPool *pool_;

        std::shared_ptr<Decryptor> GetFooterDecryptor(
            const std::string &aad,
            bool metadata);

        std::shared_ptr<Decryptor> GetColumnDecryptor(
            const std::string &column_path,
            const std::string &column_key_metadata,
            const std::string &aad,
            bool metadata = false);
    };
} // namespace kumo::pollux::parquet::arrow
